import os
import warnings
from concurrent.futures import ProcessPoolExecutor

import ants
import nibabel as nib
import numpy as np
from scipy.ndimage import gaussian_filter

warnings.filterwarnings("ignore")


class MultiAtlasRegistration:
    """多图谱配准分割类"""

    def __init__(self, atlas_dir, atlas_mask_dir, num_atlases=20):
        """
        Args:
            atlas_dir: 图谱图像目录
            atlas_mask_dir: 图谱标注目录
            num_atlases: 使用的图谱数量
        """
        self.atlas_dir = atlas_dir
        self.atlas_mask_dir = atlas_mask_dir
        self.num_atlases = num_atlases

        # 加载图谱
        self.atlas_images, self.atlas_masks = self._load_atlases()

    def _load_atlases(self):
        """加载所有图谱图像和标注"""
        atlas_images = []
        atlas_masks = []

        atlas_files = sorted(os.listdir(self.atlas_dir))[: self.num_atlases]

        print(f"加载{len(atlas_files)}个图谱...")
        for atlas_file in atlas_files:
            # 加载图像
            img_path = os.path.join(self.atlas_dir, atlas_file)
            img = nib.load(img_path)
            atlas_images.append(img.get_fdata())

            # 加载对应的标注
            mask_path = os.path.join(self.atlas_mask_dir, atlas_file)
            mask = nib.load(mask_path)
            atlas_masks.append(mask.get_fdata())

        return atlas_images, atlas_masks

    def _preprocess_image(self, image):
        """预处理图像用于配准"""
        # 强度归一化
        p1, p99 = np.percentile(image, (1, 99))
        image = np.clip(image, p1, p99)
        image = (image - np.min(image)) / (np.max(image) - np.min(image))

        # 高斯平滑去噪
        image = gaussian_filter(image, sigma=1.0)

        return image

    def _register_single_atlas(self, atlas_idx, target_image):
        """
        配准单个图谱到目标图像

        Args:
            atlas_idx: 图谱索引
            target_image: 目标图像

        Returns:
            warped_mask: 配准后的标注
        """
        # 获取图谱图像和标注
        atlas_image = self.atlas_images[atlas_idx]
        atlas_mask = self.atlas_masks[atlas_idx]

        # 预处理
        atlas_image = self._preprocess_image(atlas_image)
        target_image_prep = self._preprocess_image(target_image)

        # 转换为ANTs图像
        fixed = ants.from_numpy(target_image_prep.astype(np.float32))
        moving = ants.from_numpy(atlas_image.astype(np.float32))
        moving_mask = ants.from_numpy(atlas_mask.astype(np.float32))

        # 执行配准
        # 首先进行仿射配准
        affine_tx = ants.registration(
            fixed=fixed, moving=moving, type_of_transform="Affine", verbose=False
        )

        # 然后进行非刚性配准(SyN)
        syn_tx = ants.registration(
            fixed=fixed,
            moving=affine_tx["warpedmovout"],
            type_of_transform="SyN",
            syn_metric="CC",
            syn_sampling=2,
            reg_iterations=(100, 70, 50, 20),
            verbose=False,
        )

        # 应用变换到标注
        # 先应用仿射变换
        warped_mask = ants.apply_transforms(
            fixed=fixed,
            moving=moving_mask,
            transformlist=affine_tx["fwdtransforms"],
            interpolator="nearestNeighbor",
        )

        # 再应用非刚性变换
        warped_mask = ants.apply_transforms(
            fixed=fixed,
            moving=warped_mask,
            transformlist=syn_tx["fwdtransforms"],
            interpolator="nearestNeighbor",
        )

        return warped_mask.numpy()

    def register_all_atlases(self, target_image, parallel=True):
        """
        配准所有图谱到目标图像

        Args:
            target_image: 目标图像
            parallel: 是否并行处理

        Returns:
            warped_masks: 所有配准后的标注列表
        """
        warped_masks = []

        if parallel:
            # 并行处理
            with ProcessPoolExecutor(max_workers=4) as executor:
                futures = []
                for i in range(len(self.atlas_images)):
                    future = executor.submit(
                        self._register_single_atlas, i, target_image
                    )
                    futures.append(future)

                for i, future in enumerate(futures):
                    print(f"等待图谱 {i + 1}/{len(futures)} 配准完成...")
                    warped_mask = future.result()
                    warped_masks.append(warped_mask)
        else:
            # 串行处理
            for i in range(len(self.atlas_images)):
                print(f"配准图谱 {i + 1}/{len(self.atlas_images)}...")
                warped_mask = self._register_single_atlas(i, target_image)
                warped_masks.append(warped_mask)

        return warped_masks

    def label_fusion(
        self, warped_masks, fusion_method="majority_voting", threshold=0.5
    ):
        """
        标签融合

        Args:
            warped_masks: 配准后的标注列表
            fusion_method: 融合方法 ('majority_voting', 'weighted_voting', 'staple')
            threshold: 投票阈值

        Returns:
            fused_mask: 融合后的标注
        """
        warped_masks = np.array(warped_masks)

        if fusion_method == "majority_voting":
            # 简单多数投票
            # 计算每个位置的投票数
            votes = np.sum(warped_masks > 0, axis=0)
            # 应用阈值
            fused_mask = (votes >= threshold * len(warped_masks)).astype(np.uint8)

        elif fusion_method == "weighted_voting":
            # 加权投票（可以根据配准质量加权）
            # 这里简化为平均
            mean_prob = np.mean(warped_masks, axis=0)
            fused_mask = (mean_prob >= threshold).astype(np.uint8)

        elif fusion_method == "staple":
            # STAPLE算法 (Simultaneous Truth and Performance Level Estimation)
            fused_mask = self._staple_fusion(warped_masks)

        else:
            raise ValueError(f"Unknown fusion method: {fusion_method}")

        return fused_mask

    def _staple_fusion(self, warped_masks, max_iter=20, convergence_threshold=1e-5):
        """
        STAPLE算法实现

        Args:
            warped_masks: 配准后的标注
            max_iter: 最大迭代次数
            convergence_threshold: 收敛阈值

        Returns:
            fused_mask: 融合后的标注
        """
        num_raters = len(warped_masks)
        shape = warped_masks[0].shape

        # 初始化
        # W: 真实标注的概率估计
        W = np.mean(warped_masks, axis=0)

        # p: 敏感度 (true positive rate)
        # q: 特异度 (true negative rate)
        p = np.ones(num_raters) * 0.99
        q = np.ones(num_raters) * 0.99

        for iteration in range(max_iter):
            W_old = W.copy()

            # E步：更新W
            a = np.ones(shape)
            b = np.ones(shape)

            for j in range(num_raters):
                a *= p[j] * warped_masks[j] + (1 - p[j]) * (1 - warped_masks[j])
                b *= q[j] * (1 - warped_masks[j]) + (1 - q[j]) * warped_masks[j]

            W = a / (a + b + 1e-10)

            # M步：更新p和q
            for j in range(num_raters):
                p[j] = np.sum(W * warped_masks[j]) / (np.sum(W) + 1e-10)
                q[j] = np.sum((1 - W) * (1 - warped_masks[j])) / (np.sum(1 - W) + 1e-10)

                # 限制范围
                p[j] = np.clip(p[j], 0.5, 0.999)
                q[j] = np.clip(q[j], 0.5, 0.999)

            # 检查收敛
            if np.mean(np.abs(W - W_old)) < convergence_threshold:
                print(f"STAPLE收敛于第{iteration}次迭代")
                break

        # 二值化
        fused_mask = (W >= 0.5).astype(np.uint8)

        return fused_mask

    def segment(self, target_image_path, output_path=None):
        """
        对目标图像进行分割

        Args:
            target_image_path: 目标图像路径
            output_path: 输出路径（可选）

        Returns:
            segmentation: 分割结果
            metrics: 性能指标（如果有真实标注）
        """
        # 加载目标图像
        target_nib = nib.load(target_image_path)
        target_image = target_nib.get_fdata()

        print("开始多图谱配准分割...")

        # 配准所有图谱
        warped_masks = self.register_all_atlases(target_image)

        # 标签融合
        print("执行标签融合...")
        segmentation = self.label_fusion(warped_masks, fusion_method="majority_voting")

        # 保存结果
        if output_path:
            seg_nib = nib.Nifti1Image(segmentation, target_nib.affine)
            nib.save(seg_nib, output_path)
            print(f"分割结果保存至: {output_path}")

        return segmentation


class MultiAtlasWithDL:
    """
    结合深度学习的多图谱配准
    使用深度学习模型改进配准结果
    """

    def __init__(self, atlas_registration, dl_model):
        self.atlas_registration = atlas_registration
        self.dl_model = dl_model

    def hybrid_segment(self, target_image_path):
        """
        混合分割方法

        Args:
            target_image_path: 目标图像路径

        Returns:
            final_segmentation: 最终分割结果
        """
        # 1. 多图谱配准分割
        atlas_seg = self.atlas_registration.segment(target_image_path)

        # 2. 深度学习分割
        # 这里需要调用深度学习模型
        # dl_seg = self.dl_model.predict(target_image_path)

        # 3. 融合两种方法的结果
        # final_segmentation = self._fuse_results(atlas_seg, dl_seg)

        # 暂时返回图谱分割结果
        return atlas_seg